"""Model convert

Convert PyTorch model to OM model by end-to-end.

Example:
    >>> python tools/convert.py \
    >>>     mmdetection/configs/yolox/yolox_s_8xb8-300e_coco.py \
    >>>     models/yolox_s_8x8_300e_coco_20211121_095711-4592a793.pth \
    >>>     --batch_sizes 1 \
    >>>     --input_shape 640 640 \
    >>>     --save_dir models \
    >>>     --file_prefix yolox_s
"""

import argparse
from pathlib import Path
from npu_mmdet.convert import ModelConverter


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
                        description='convert pytorch model to OM model.')
    parser.add_argument('config', type=str, help='path to model config')
    parser.add_argument('checkpoint', type=str,
                        help='path to model checkpoint.')
    parser.add_argument('--task', type=str,
                        default='det', choices=['det', 'seg'],
                        help='object detection or instance segmentation?')
    parser.add_argument('--image', type=str,
                        default='./demo/test_images/demo.jpg',
                        help='image used to convert model.')
    parser.add_argument('--batch_sizes', nargs='+', type=int,
                        help='specify one or more batch size values.')
    parser.add_argument('--use_dynamic_batch_size', action='store_true',
                        help='a dynamic batch size OM will be generated if True'
                             ', otherwise multiple static OM.')
    parser.add_argument('--input_shape', nargs=2, type=int,
                        default=[800, 1088],
                        help='height and width of input data.')
    parser.add_argument('--soc_version', type=str,
                        help='version of the target SoC. '
                             'eg: `--soc_version Ascend310P3`')
    parser.add_argument('--save_dir', type=str,
                        default=None,
                        help='the dir to save models.')
    parser.add_argument('--file_prefix', type=str,
                        default='model',
                        help='the file prefix to save models.')

    args = parser.parse_args()

    if not args.save_dir:
        args.save_dir = Path(args.config).stem
    save_dir = Path(args.save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)
    save_prefix = str(save_dir / args.file_prefix)

    # ==================== conventional way =========================
    # converter = ModelConverter(args.config,
    #                            task=args.task,
    #                            input_shape=args.input_shape)
    # converter.convert(args.checkpoint,
    #                   input_image=args.image,
    #                   save_prefix=save_prefix,
    #                   batch_sizes=args.batch_sizes,
    #                   soc_version=args.soc_version,
    #                   use_dynamic_batch_size=args.use_dynamic_batch_size)

    # ========================== fast way ===========================
    ModelConverter.fast_convert(
        args.config, args.checkpoint,
        task=args.task,
        input_shape=args.input_shape,
        input_image=args.image,
        save_prefix=save_prefix,
        batch_sizes=args.batch_sizes,
        soc_version=args.soc_version,
        use_dynamic_batch_size=args.use_dynamic_batch_size
    )